#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
绘图模块

该模块提供了用于可视化路径规划过程的Plotter类。
"""

import os
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Circle, Rectangle, Ellipse


class Plotter:
    """
    用于可视化路径规划过程的类。
    """
    
    def __init__(self, map_instance, robot, save_dir=None, config=None):
        """
        使用地图和机器人初始化绘图器。
        
        参数:
            map_instance: 表示环境的地图对象
            robot: 表示机器人的对象
            save_dir (str, optional): 保存动画帧的目录
            config (dict, optional): 配置字典
        """
        self.map = map_instance
        self.robot = robot
        self.save_dir = save_dir
        self.fig = None
        self.ax = None
        self.frame_count = 0
        self.config = config or {}
        
        # 提取可视化配置
        self.viz_config = self.config.get("visualization", {})
        self.show_narrow_passages = self.viz_config.get("show_narrow_passages", True)
        
        # 提取颜色
        colors = self.viz_config.get("colors", {})
        self.tree_color = colors.get("tree", [0.7, 0.7, 0.7, 0.5])
        self.path_color = colors.get("path", [0.0, 0.8, 0.2, 1.0])
        self.start_color = colors.get("start", [0.0, 1.0, 0.0, 1.0])
        self.goal_color = colors.get("goal", [1.0, 0.0, 0.0, 1.0])
        self.obstacle_color = colors.get("obstacles", [0.1, 0.1, 0.1, 1.0])
        self.narrow_passage_color = colors.get("narrow_passages", [1.0, 0.7, 0.0, 0.5])
        
        # 如果指定了save_dir，则初始化绘图
        if save_dir and not os.path.exists(save_dir):
            os.makedirs(save_dir, exist_ok=True)
    
    def init_planning(self, tree, start, goal):
        """
        初始化规划过程的绘图。
        
        参数:
            tree: RRT树
            start (tuple): 起始位置 (x, y)
            goal (tuple): 目标位置 (x, y)
        """
        # 如果尚未创建，则创建图形和坐标轴
        if self.fig is None or self.ax is None:
            # 如果指定了绘图分辨率，则获取
            plot_resolution = self.viz_config.get("plot_resolution", [1000, 800])
            figsize = (plot_resolution[0]/100, plot_resolution[1]/100)
            self.fig, self.ax = plt.subplots(figsize=figsize)
            
        # 根据地图尺寸设置绘图范围
        width, height = self.map.get_dimensions()
        self.ax.set_xlim([-width * 0.05, width * 1.05])
        self.ax.set_ylim([-height * 0.05, height * 1.05])
        
        # 绘制地图边界
        self.ax.add_patch(Rectangle((0, 0), width, height, fill=False, edgecolor='black'))
        
        # 绘制障碍物
        self._plot_obstacles()
        
        # 用较大的标记绘制起点和终点，并加上标签
        self.ax.plot(start[0], start[1], 'o', markersize=15, label='Start', 
                    color=self.start_color)
        self.ax.plot(goal[0], goal[1], 'o', markersize=15, label='Goal', 
                    color=self.goal_color)
        
        # 在点旁边添加文本标签
        self.ax.text(start[0]+5, start[1]+5, "Start", fontsize=12)
        self.ax.text(goal[0]+5, goal[1]+5, "Goal", fontsize=12)
        
        # 添加图例和网格
        self.ax.legend()
        self.ax.grid(True)
        self.ax.set_title('RRT路径规划')
        self.ax.set_xlabel('X')
        self.ax.set_ylabel('Y')
        
        # 如果指定了save_dir，则保存初始帧
        if self.save_dir:
            self._save_frame()
    
    def init_bidirectional_planning(self, start_tree, goal_tree, start, goal):
        """
        初始化双向规划的绘图。
        
        参数:
            start_tree: 起始RRT树
            goal_tree: 目标RRT树
            start (tuple): 起始位置 (x, y)
            goal (tuple): 目标位置 (x, y)
        """
        # 初始化基本绘图
        self.init_planning(start_tree, start, goal)
        
        # 更新双向规划的标题
        self.ax.set_title('双向RRT路径规划')
    
    def update_planning(self, tree, nearest_node_to_goal, goal_node=None, path=None, narrow_passages=None):
        """
        在规划过程中更新绘图。
        
        参数:
            tree: RRT树
            nearest_node_to_goal: 最接近目标的节点
            goal_node (optional): 如果达到，则为目标节点
            path (optional): 当前最佳路径
            narrow_passages (optional): 检测到的狭窄通道列表
        """
        if self.fig is None or self.ax is None:
            return
            
        # 清除之前的RRT树
        for artist in self.ax.lines[:]:
            if artist.get_label() not in ['Start', 'Goal', 'Path']:
                artist.remove()
                
        # 清除之前的狭窄通道
        for patch in self.ax.patches[:]:
            if getattr(patch, "_is_narrow_passage", False):
                patch.remove()
        
        # 绘制RRT树边
        edges = tree.get_edges()
        for edge in edges:
            (x1, y1), (x2, y2) = edge
            self.ax.plot([x1, x2], [y1, y2], '-', linewidth=0.5, alpha=0.5, 
                        color=self.tree_color)
            
        # 绘制RRT树节点
        xs, ys = tree.get_node_positions()
        self.ax.plot(xs, ys, '.', markersize=2, color=self.tree_color)
        
        # 如果有可用的路径，则绘制最佳路径
        if path:
            path_xs = [point[0] for point in path]
            path_ys = [point[1] for point in path]
            for artist in self.ax.lines[:]:
                if artist.get_label() == 'Path':
                    artist.remove()
            self.ax.plot(path_xs, path_ys, '-', linewidth=2, label='Path', 
                        color=self.path_color)
        
        # 绘制最接近目标的节点
        if nearest_node_to_goal:
            self.ax.plot(nearest_node_to_goal.x, nearest_node_to_goal.y, '*', 
                        markersize=10, color=[1.0, 1.0, 0.0, 1.0])
            
        # 如果达到目标节点，则绘制
        if goal_node:
            self.ax.plot(goal_node.x, goal_node.y, '*', markersize=10, 
                        color=self.goal_color)
        
        # 如果有可用的狭窄通道并启用，则绘制
        if narrow_passages and self.show_narrow_passages:
            self._plot_narrow_passages(narrow_passages)
            
        # 绘制更新的图
        self.fig.canvas.draw()
        plt.pause(0.001)  # 短暂暂停以允许显示更新
        
        # 如果指定了save_dir，则保存帧
        if self.save_dir:
            self._save_frame()
    
    def update_bidirectional_planning(self, start_tree, goal_tree, 
                                     connection_node_start=None, 
                                     connection_node_goal=None, path=None,
                                     narrow_passages=None):
        """
        在双向规划过程中更新绘图。
        
        参数:
            start_tree: 起始RRT树
            goal_tree: 目标RRT树
            connection_node_start (optional): 起始树中的连接节点
            connection_node_goal (optional): 目标树中的连接节点
            path (optional): 当前最佳路径
            narrow_passages (optional): 检测到的狭窄通道列表
        """
        if self.fig is None or self.ax is None:
            return
            
        # 清除之前的RRT树
        for artist in self.ax.lines[:]:
            if artist.get_label() not in ['Start', 'Goal', 'Path']:
                artist.remove()
                
        # 清除之前的狭窄通道
        for patch in self.ax.patches[:]:
            if getattr(patch, "_is_narrow_passage", False):
                patch.remove()
                
        # 绘制起始树边（蓝色）
        edges = start_tree.get_edges()
        for edge in edges:
            (x1, y1), (x2, y2) = edge
            self.ax.plot([x1, x2], [y1, y2], '-', linewidth=0.5, alpha=0.5, 
                        color=[0.0, 0.0, 1.0, 0.5])
            
        # 绘制起始树节点
        xs, ys = start_tree.get_node_positions()
        self.ax.plot(xs, ys, '.', markersize=2, color=[0.0, 0.0, 1.0, 0.8])
        
        # 绘制目标树边（绿色）
        edges = goal_tree.get_edges()
        for edge in edges:
            (x1, y1), (x2, y2) = edge
            self.ax.plot([x1, x2], [y1, y2], '-', linewidth=0.5, alpha=0.5, 
                        color=[0.0, 0.8, 0.0, 0.5])
            
        # 绘制目标树节点
        xs, ys = goal_tree.get_node_positions()
        self.ax.plot(xs, ys, '.', markersize=2, color=[0.0, 0.8, 0.0, 0.8])
        
        # 如果有可用的树之间的连接，则绘制
        if connection_node_start and connection_node_goal:
            self.ax.plot([connection_node_start.x, connection_node_goal.x], 
                        [connection_node_start.y, connection_node_goal.y], 
                        '-', linewidth=2, color=[1.0, 1.0, 0.0, 1.0])
            
            # 突出显示连接节点
            self.ax.plot(connection_node_start.x, connection_node_start.y, '*', 
                        markersize=10, color=[1.0, 1.0, 0.0, 1.0])
            self.ax.plot(connection_node_goal.x, connection_node_goal.y, '*', 
                        markersize=10, color=[1.0, 1.0, 0.0, 1.0])
        
        # 如果有可用的路径，则绘制最佳路径
        if path:
            path_xs = [point[0] for point in path]
            path_ys = [point[1] for point in path]
            for artist in self.ax.lines[:]:
                if artist.get_label() == 'Path':
                    artist.remove()
            self.ax.plot(path_xs, path_ys, '-', linewidth=2, label='Path', 
                        color=self.path_color)
        
        # 如果有可用的狭窄通道并启用，则绘制
        if narrow_passages and self.show_narrow_passages:
            self._plot_narrow_passages(narrow_passages)
        
        # 绘制更新的图
        self.fig.canvas.draw()
        plt.pause(0.001)  # 短暂暂停以允许显示更新
        
        # 如果指定了save_dir，则保存帧
        if self.save_dir:
            self._save_frame()
    
    def plot_result(self, path, nodes=None, narrow_passages=None):
        """
        绘制最终规划结果。
        
        参数:
            path: 作为(x, y)坐标列表的最终路径
            nodes (optional): 要突出显示的节点列表
            narrow_passages (optional): 检测到的狭窄通道列表
        """
        if self.fig is None or self.ax is None:
            return
            
        # 清除之前的路径
        for artist in self.ax.lines[:]:
            if artist.get_label() == 'Path':
                artist.remove()
                
        # 清除之前的狭窄通道
        for patch in self.ax.patches[:]:
            if getattr(patch, "_is_narrow_passage", False):
                patch.remove()
                
        # 绘制最终路径
        if path:
            path_xs = [point[0] for point in path]
            path_ys = [point[1] for point in path]
            self.ax.plot(path_xs, path_ys, '-', linewidth=2, label='Path', 
                        color=self.path_color)
            
            # 突出显示路径点
            self.ax.plot(path_xs, path_ys, '.', markersize=5, 
                        color=self.path_color)
        
        # 如果提供了特定节点，则突出显示
        if nodes:
            for node in nodes:
                self.ax.plot(node.x, node.y, '*', markersize=6, 
                           color=[1.0, 1.0, 0.0, 1.0])
        
        # 如果有可用的狭窄通道并启用，则绘制
        if narrow_passages and self.show_narrow_passages:
            self._plot_narrow_passages(narrow_passages)
                
        # 更新图例
        self.ax.legend()
        
        # 绘制更新的图
        self.fig.canvas.draw()
        
        # 如果指定了save_dir，则保存最终帧
        if self.save_dir:
            self._save_frame(final=True)
    
    def show(self):
        """
        显示绘图。
        """
        if self.fig:
            plt.show()
    
    def save_figure(self, filename):
        """
        将绘图保存到文件。
        
        参数:
            filename (str): 输出文件名
        """
        if self.fig:
            self.fig.savefig(filename, dpi=300, bbox_inches='tight')
    
    def close(self):
        """
        关闭绘图。
        """
        if self.fig:
            plt.close(self.fig)
            self.fig = None
            self.ax = None
    
    def _plot_obstacles(self):
        """
        绘制地图中的障碍物。
        """
        obstacles = self.map.get_obstacles()
        for obstacle in obstacles:
            # 获取障碍物边界用于绘图
            xs, ys = obstacle.get_boundaries()
            
            if obstacle.type == "circle":
                self.ax.add_patch(Circle((obstacle.x, obstacle.y), 
                                        obstacle.radius, 
                                        color=self.obstacle_color))
            elif obstacle.type == "rectangle":
                self.ax.add_patch(Rectangle((obstacle.x, obstacle.y), 
                                          obstacle.width, obstacle.height, 
                                          color=self.obstacle_color))
            else:
                # 对于其他障碍物类型，使用返回的边界
                self.ax.fill(xs, ys, color=self.obstacle_color)
    
    def _plot_narrow_passages(self, narrow_passages):
        """
        绘制检测到的狭窄通道。
        
        参数:
            narrow_passages: 狭窄通道区域列表 (x, y, radius)
        """
        if not narrow_passages:
            return
            
        for passage in narrow_passages:
            x, y, radius = passage
            circle = Circle((x, y), radius, color=self.narrow_passage_color, alpha=0.6)
            circle._is_narrow_passage = True
            self.ax.add_patch(circle)
            
            # 在狭窄通道上添加标签
            self.ax.text(x, y, "NP", fontsize=8, 
                       horizontalalignment='center', verticalalignment='center')
    
    def _save_frame(self, final=False):
        """
        将当前绘图保存为动画帧。
        
        参数:
            final (bool): 是否为最终帧
        """
        if not self.save_dir:
            return
            
        # 创建帧目录
        frames_dir = os.path.join(self.save_dir, 'frames')
        if not os.path.exists(frames_dir):
            os.makedirs(frames_dir, exist_ok=True)
            
        # 保存帧
        if final:
            filename = os.path.join(self.save_dir, 'final_path.png')
            self.fig.savefig(filename, dpi=300, bbox_inches='tight')
        else:
            filename = os.path.join(frames_dir, f'frame_{self.frame_count:04d}.png')
            self.fig.savefig(filename, dpi=150, bbox_inches='tight')
            self.frame_count += 1 